package main

import "fmt"

/*
给定整数数组 nums 和整数 k，请返回数组中第 k 个最大的元素。

请注意，你需要找的是数组排序后的第 k 个最大的元素，而不是第 k 个不同的元素。

你必须设计并实现时间复杂度为 O(n) 的算法解决此问题。
*/

// 先完整快排（首元素为 pivot，不做随机），再取第 k 大
func findKthLargestV1(nums []int, k int) int {
	n := len(nums)
	if k < 1 || k > n {
		panic("k out of range")
	}
	quickSortV1(nums, 0, n-1)
	return nums[n-k]
}

// Hoare 风格（首元素 pivot）
func quickSortV1(a []int, l, r int) {
	if l >= r { // 剩一个元素不用排
		return
	}
	pivot := a[l] // 去最左端元素做pivot
	i, j := l, r
	for i < j {
		// 从右往左找比privot大的
		for i < j && a[j] >= pivot {
			j--
		}
		// 从左往右找比privot小的
		for i < j && a[i] <= pivot {
			i++
		}
		if i < j { // 必须要比较再交换，否则会存在i>j的越界情况，i=j交换没有意义
			a[i], a[j] = a[j], a[i]
		}
	}
	// i == j，放置 pivot
	a[l], a[i] = a[i], a[l]
	quickSortV1(a, l, i-1)
	quickSortV1(a, i+1, r)
}

// 优化，根据第k大的性质，只需要判断pivot是否为第k大的元素（根据下标位置），否则只递归左边或右边，可以在
func findKthLargestV2(nums []int, k int) int {
	n := len(nums)
	if k < 1 || k > n {
		panic("k out of range")
	}
	target := n - k // 第 k 大 -> 升序下标 target
	return quickSelect(nums, 0, n-1, target)
}

func quickSelect(a []int, l, r, target int) int {
	for l <= r {
		if l == r {
			return a[l]
		}
		pivot := a[l]
		i, j := l, r
		for i < j {
			for i < j && a[j] >= pivot {
				j--
			}
			for i < j && a[i] <= pivot {
				i++
			}
			if i < j {
				a[i], a[j] = a[j], a[i]
			}
		}
		// 放 pivot
		a[l], a[i] = a[i], a[l]

		if i == target {
			return a[i]
		} else if i < target {
			l = i + 1
		} else {
			r = i - 1
		}
	}
	return -1
}

func main() {
	nums := []int{3, 2, 1, 5, 6, 4}
	k := 2
	fmt.Println(findKthLargestV2(nums, k)) // 5
}
